import backtrader as bt
import akshare as ak
import pandas as pd
from datetime import date, datetime, timedelta
from select_stock import get_stocks_sector, last_day_of_month, run_screener
from data_feed import StockDataFetcher, MyStockData

# ====选股日期=======
year = 2024
month = 6
symbol="000300"
num=10  # 阳线根数
head=10 # 符合条件前十名

# ====策略时间=============
from_date = '2024,6, 20'
to_date = '2024,7, 30'

# ====选股日期=======
# today = datetime.now().date()
today=last_day_of_month(year, month)
today_str = today.strftime("%Y%m%d")
today = datetime.strptime(today_str, "%Y%m%d")
start_date=(today - timedelta(days=30)).strftime("%Y%m%d") 
end_date=today.strftime("%Y%m%d")
print(f"开始日期: {start_date}")
# 获取板块股票
stock_codes= get_stocks_sector(symbol)
# 选出阳线大于几根的股票
select_stock=run_screener(stock_codes=stock_codes, 
                            start_date=start_date, 
                            end_date=end_date, 
                            num=num,
                            head=head)
codes=select_stock['股票代码'].tolist()
print("选出%d个股票列表: %s" % (num, select_stock))

# =========================================================================
class MultiTestStrategy(bt.Strategy):
    params = (
        ("pfast", 5),  # period for the fast moving average
        ("pslow",10),   # period for the slow moving average
        ("position_limit",10),    #限制持仓数量
    )

    def prenext(self):
        pass

    def downcast(self, amount, lot):
        return abs(amount // lot * lot)

    def __init__(self):
        # 止损比例
        self.stop_loss=0.00
        # 记录买入价格
        self.buy_price = None
        # 纪录交易执行方向, 打印输出
        self.ordersta=''        # 纪录交易执行方向
        self.print_list = []    # 储存交易信息
        self.now_lst = []
        self.order = None
        self.buy_list = []  # 已经购买的股票代码列表
        # 添加移动平均线指标，循环计算每个股票的指标
        # self.getdatanames()：这是一个方法调用，返回当前策略中所有可用的数据名称（即股票代码）的列表。
        # self.sma = {x: bt.ind.SMA(self.getdatabyname(x), period=self.p.pslow) for x in self.getdatanames()}
        self.sma1 = {x: bt.ind.SMA(self.getdatabyname(x), period=self.p.pfast) for x in self.getdatanames()}
        self.sma2 = {x: bt.ind.SMA(self.getdatabyname(x), period=self.p.pslow) for x in self.getdatanames()}
        self.crossover = {x: bt.ind.CrossOver(self.sma1[x], self.sma2[x]) for x in self.getdatanames()} 
        
    def next(self):
        if len(self.buy_list) < self.p.position_limit:# 限制持仓数量
            for secu in set(self.getdatanames()) - set(self.buy_list):#遍历所有尚未购买的股票代码，以便在这些股票上执行相应的操作
                data = self.getdatabyname(secu)
                # 当前交易日的日期（data.datetime.date(1)）小于预定的结束日期（end_dates[secu]）
                if self.crossover[secu] > 0 and data.datetime.date(-1) < end_dates[secu]:
                    self.ordersta='买入'
                    order_value = self.broker.getvalue() * 0.18
                    order_amount = self.downcast(order_value / data.close[0], 100)#data=self.getpositionbyname(x)
                    self.order = self.buy(data, size=order_amount, name=secu)
                    self.log(f"买{secu}, price:{data.close[0]:.2f}, amout: {order_amount}")
                    self.buy_list.append(secu)
                    
        if self.buy_list :
            self.now_lst = self.buy_list.copy()  # 创建一个当前持仓列表的副本
            for secu in self.now_lst:
                data = self.getdatabyname(secu)
                # if data.close < self.sma[secu] or data.datetime.date(1) >= end_dates[secu]:
                if self.crossover[secu] < 0 or data.datetime.date(-1) == end_dates[secu]:
                    self.ordersta=('卖出') 
                    # 0：这是目标持仓比例，表示希望该股票的持仓比例达到 0%，即完全卖出。
                    self.order = self.order_target_percent(data, 0, name=secu)
                    self.log(f"卖{secu}, price:{data.close[0]:.2f}, pct: 0")
                    # continue
                    self.now_lst.remove(secu)# 从副本中移除满足卖出条件的股票代码
                    # self.now_lst.append(secu)
                self.buy_list = self.now_lst# 更新当前持仓列表

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed, order.Canceled, order.Margin]:
            if order.isbuy():
                self.log(f"""买入{order.info['name']}, 成交量{order.executed.size}，成交价{order.executed.price:.2f}""")
                self.log(
                    f'资产：{self.broker.getvalue():.2f} 持仓：{[(x, self.getpositionbyname(x).size) for x in self.buy_list]}')
            elif order.issell():
                self.log(f"""卖出{order.info['name']}, 成交量{order.executed.size}，成交价{order.executed.price:.2f}""")
                self.log(
                    f'资产：{self.broker.getvalue():.2f} 持仓：{[(x, self.getpositionbyname(x).size) for x in self.buy_list]}')
            self.bar_executed = len(self)
        self.order = None

    def log(self, txt, dt=None):
        dt = dt or self.datas[0].datetime.date(0)
        print('%s , %s' % (dt.isoformat(), txt))
    
# ======================================================================
from_date = datetime(from_date)
to_date = datetime(to_date)
cerebro = bt.Cerebro()

end_dates = {} #用于存储每个股票代码（code）对应的数据结束日期。确保数据完整性
for code in codes:
    stock_fetcher = StockDataFetcher(symbol=code)
    data = stock_fetcher.get_data()
    data_feed = MyStockData(dataname=data, 
                            fromdate=from_date, todate=to_date)
    
    cerebro.adddata(data_feed, name=code)
    # end_dates[code] = data.index[-1] 
    # 这行代码从data（假设它是一个Pandas DataFrame）中获取最后一条记录的索引
    # （即日期），并将其作为结束日期存储在end_dates字典中，键为股票代码。
    end_dates[code] = data.index[-1]# 当策略放到库文件,策略引用end_dates需要传参数
    print('添加股票数据：code: %s' % code)

cerebro.broker.setcash(1000000.0)
cerebro.broker.setcommission(commission=0.001)
cerebro.addstrategy(MultiTestStrategy)
# cerebro.addstrategy(MultiSmaCross, end_dates=end_dates)
cerebro.run()

portvalue = cerebro.broker.getvalue()
print(f'结束资金: {round(portvalue, 2)}')
cerebro.plot()